package com.aiwiown.snackmq.cluster.raft.manager;

import com.aiwiown.snackmq.cluster.config.ClusterConfig;
import com.aiwiown.snackmq.cluster.raft.SnackMQRaftNodeConfig;
import com.aiwiown.snackmq.cluster.raft.SnackMQStateMachine;
import com.aiwiown.snackmq.cluster.raft.entity.SnackMQRaftLogEntry;
import com.alipay.sofa.jraft.Closure;
import com.alipay.sofa.jraft.Node;
import com.alipay.sofa.jraft.RaftGroupService;
import com.alipay.sofa.jraft.StateMachine;
import com.alipay.sofa.jraft.Status;
import com.alipay.sofa.jraft.conf.Configuration;
import com.alipay.sofa.jraft.entity.PeerId;
import com.alipay.sofa.jraft.entity.Task;
import com.alipay.sofa.jraft.entity.UserLog;
import com.alipay.sofa.jraft.error.RaftError;
import com.alipay.sofa.jraft.option.NodeOptions;
import com.alipay.sofa.jraft.option.RaftOptions;
import com.alipay.sofa.jraft.rpc.RpcServer;
import com.alipay.sofa.jraft.util.Endpoint;
import lombok.extern.slf4j.Slf4j;

import java.io.File;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * 【已修复】Raft 节点管理器。
 * 封装了 JRaft Node 的生命周期管理、日志应用和 Leader 信息查询。
 * 这是确保集群中每个 Raft 组（元数据、数据）正确运行的核心组件。
 */
@Slf4j
public class SnackMQRaftNodeManager {

    private final RaftGroupService raftGroupService;
    private Node node;
    private final AtomicBoolean started = new AtomicBoolean(false);
    private final ClusterConfig clusterConfig;

    /**
     * 构造函数，用于初始化 Raft 节点管理器
     * @param config Raft 节点配置
     * @param fsm 状态机，用于处理 Raft 日志条目
     * @param rpcServer RPC 服务器，用于处理 Raft 节点间的通信
     * @param clusterConfig 集群配置，包含了 Raft 地址到业务地址的映射
     */
    public SnackMQRaftNodeManager(SnackMQRaftNodeConfig config, StateMachine fsm, RpcServer rpcServer, ClusterConfig clusterConfig) {
        this.clusterConfig = clusterConfig;
        try {
            NodeOptions nodeOptions = createNodeOptions(config, fsm);
            this.raftGroupService = new RaftGroupService(config.getGroupId(), new PeerId(config.getHost(), config.getRaftPort()), nodeOptions, rpcServer, true);
            // --- 【核心修复 1】: 移除此处的 getNode() 调用。Node 实例将在 start() 方法中被正确赋值。
            // this.node = this.raftGroupService.getNode(); // <-- 移除此行
        } catch (IOException e) {
            throw new RuntimeException("Failed to initialize Raft node manager", e);
        }
    }

    private NodeOptions createNodeOptions(SnackMQRaftNodeConfig config, StateMachine fsm) throws IOException {
        NodeOptions nodeOptions = new NodeOptions();
        nodeOptions.setElectionTimeoutMs(config.getElectionTimeoutMs());
        nodeOptions.setFsm(fsm);
        // --- 【核心修复 2】: 此处现在可以正常编译，因为 SnackMQRaftNodeConfig 已包含此方法。
        nodeOptions.setSnapshotIntervalSecs(config.getSnapshotIntervalSecs());
        nodeOptions.setSnapshotIntervalSecs(300);
        // --- 【核心性能优化】: 获取内部的 RaftOptions 并进行配置 ---
        final RaftOptions raftOptions = nodeOptions.getRaftOptions();
        // 默认情况下，JRaft 对每个日志条目都执行 fsync，这在高性能场景下会成为瓶颈。
        // 通过禁用单次同步并设置缓冲区，我们可以利用“组提交”来大幅提高吞吐量。
        // 关闭同步刷盘，由 JRaft 内部的 Group-Commit 机制来保证数据持久化
        raftOptions.setSync(false);
        // 设置 Raft 日志的写缓冲区大小，例如 512KB。当缓冲区满或超时时，会批量刷盘。
        raftOptions.setMaxAppendBufferSize(512 * 1024);

        File raftDataDir = new File(config.getRaftDataPath());
        if (!raftDataDir.exists() && !raftDataDir.mkdirs()) {
            throw new IOException("无法创建 Raft 数据目录: " + raftDataDir.getAbsolutePath());
        }

        nodeOptions.setLogUri(config.getRaftDataPath() + File.separator + "log");
        nodeOptions.setRaftMetaUri(config.getRaftDataPath() + File.separator + "meta");
        nodeOptions.setSnapshotUri(config.getRaftDataPath() + File.separator + "snapshot");

        Configuration initialConf = new Configuration();
        if (!initialConf.parse(config.getInitialConf())) {
            throw new IOException("Failed to parse initial Raft configuration: " + config.getInitialConf());
        }
        nodeOptions.setInitialConf(initialConf);
        return nodeOptions;
    }

//    private NodeOptions createNodeOptions(SnackMQRaftNodeConfig config, StateMachine fsm) throws IOException {
//        NodeOptions nodeOptions = new NodeOptions();
//        nodeOptions.setElectionTimeoutMs(config.getElectionTimeoutMs());
//        nodeOptions.setFsm(fsm);
//        // --- 【核心修复 2】: 此处现在可以正常编译，因为 SnackMQRaftNodeConfig 已包含此方法。
//        nodeOptions.setSnapshotIntervalSecs(config.getSnapshotIntervalSecs());
//
//        File raftDataDir = new File(config.getRaftDataPath());
//        if (!raftDataDir.exists() && !raftDataDir.mkdirs()) {
//            throw new IOException("无法创建 Raft 数据目录: " + raftDataDir.getAbsolutePath());
//        }
//
//        nodeOptions.setLogUri(config.getRaftDataPath() + File.separator + "log");
//        nodeOptions.setRaftMetaUri(config.getRaftDataPath() + File.separator + "meta");
//        nodeOptions.setSnapshotUri(config.getRaftDataPath() + File.separator + "snapshot");
//
//        Configuration initialConf = new Configuration();
//        if (!initialConf.parse(config.getInitialConf())) {
//            throw new IOException("Failed to parse initial Raft configuration: " + config.getInitialConf());
//        }
//        nodeOptions.setInitialConf(initialConf);
//        return nodeOptions;
//    }

    public boolean start(boolean startRpcServer) {
        if (started.compareAndSet(false, true)) {
            try {
                // 在新版本 JRaft 中，start() 返回 Node，所以我们更新 this.node
                this.node = this.raftGroupService.start(startRpcServer);
                if (this.node == null) {
                    throw new RuntimeException("Failed to start Raft node, node is null.");
                }
                log.info("Raft node started successfully. Node status: <{}>", this.node.getNodeId());
                return true;
            } catch (Exception e) {
                log.error("Failed to start Raft node", e);
                started.set(false);
                return false;
            }
        }
        return true;
    }

    public void shutdown() {
        if (started.compareAndSet(true, false)) {
            log.info("Shutting down Raft node manager for group {}...", raftGroupService.getGroupId());
            try {
                this.raftGroupService.shutdown();
                log.info("Raft group service for group {} shutdown completed.", raftGroupService.getGroupId());
            } catch (Exception e) {
                log.error("Error shutting down Raft group service for group {}", raftGroupService.getGroupId(), e);
            }
            log.info("Raft node manager for group {} shutdown successfully.", raftGroupService.getGroupId());
        }
    }

    public void apply(SnackMQRaftLogEntry logEntry, Closure closure) {
        if (this.node == null) {
            if (closure != null) {
                // 如果节点未启动，返回一个明确的错误
                closure.run(new Status(RaftError.EPERM, "Raft node is not started."));
            }
            return;
        }

        // --- 【最终修复】: 移除 isLeader() 检查，将任务直接交给 JRaft Node 处理。---
        // JRaft Node 内部会自动处理请求转发。如果当前节点不是 Leader，它会将任务转发给 Leader。
        // 如果没有 Leader，它会返回一个包含正确错误码的 Status，由 Closure 处理。
        // 之前的 isLeader() 检查是多余的，并且它错误地短路了 JRaft 的内置重定向/转发机制。
        final Task task = new Task(ByteBuffer.wrap(logEntry.toBytes()), closure);
        this.node.apply(task);
    }

    /**
     * 【新增】根据索引从 Raft 日志中读取单个日志条目。
     * 这是延迟消息恢复功能的关键，它允许 TimerScannerService 恢复原始消息。
     *
     * @param index Raft 日志的索引 (即 commitLogOffset)
     * @return 反序列化后的 SnackMQRaftLogEntry，如果找不到或出错则返回 null
     */
    public SnackMQRaftLogEntry readLogEntry(long index) {
        if (this.node == null) {
            log.error("Cannot read log entry, Raft node is not initialized.");
            return null;
        }
        try {
            // 【修复】将接收变量的类型从 LogEntry 修改为 UserLog
            UserLog userLog = this.node.readCommittedUserLog(index);
            if (userLog == null) {
                log.warn("Log entry not found at index {}.", index);
                return null;
            }
            // 从 UserLog 中获取数据并反序列化为我们自己的类型
            return SnackMQRaftLogEntry.fromBytes(userLog.getData().array());
        } catch (Exception e) {
            log.error("Failed to read log entry at index {}.", index, e);
            return null;
        }
    }

    public boolean isLeader() {
        return this.node != null && this.node.isLeader();
    }

    /**
     * --- 【核心修复】 ---
     * 获取当前 Leader 的业务地址 (host:business_port)。
     * 此方法现在完全依赖于从 BrokerConfig 传递下来的 `raftAddressToBusinessPortMap` 来进行地址解析。
     *
     * @return Leader 的业务地址字符串，如果找不到 Leader 或无法解析地址，则返回 null。
     */
    public String getLeaderAddress() {
        if (this.node == null) {
            return null;
        }
        final PeerId leader = this.node.getLeaderId();
        if (leader == null || leader.isEmpty()) {
            log.warn("No leader found for Raft group {}.", getGroupId());
            return null;
        }

        // 从 ClusterConfig 中获取 Raft 地址到业务端口的映射
        Map<String, Integer> addressMap = clusterConfig.getRaftAddressToBusinessPortMap();
        if (addressMap == null || addressMap.isEmpty()) {
            log.error("CRITICAL: Raft address to business port map is not configured. Cannot provide redirect address.");
            return null;
        }

        String leaderRaftAddress = leader.getIp() + ":" + leader.getPort();
        Integer businessPort = addressMap.get(leaderRaftAddress);

        if (businessPort == null) {
            log.error("Could not resolve business port for leader Raft address {}. Cannot provide redirect address. Available map: {}",
                    leaderRaftAddress, addressMap);
            return null;
        }

        return leader.getIp() + ":" + businessPort;
    }

    public List<PeerId> getAlivePeers() {
        if (this.node != null && this.node.isLeader()) {
            return this.node.listAlivePeers();
        }
        return new ArrayList<>();
    }

    public String getGroupId() {
        return this.node != null ? this.node.getGroupId() : "unknown";
    }

    public PeerId getSelfPeerId() {
        return this.node != null ? this.node.getNodeId().getPeerId() : null;
    }

    /**
     * 【修复二】: 新增 getLeaderEndpoint() 方法。
     * 这是实现重定向逻辑的关键，它提供了 Leader 的原始 Raft 地址。
     *
     * @return Leader 的 Endpoint 对象，如果找不到 Leader 则返回 null。
     */
    public Endpoint getLeaderEndpoint() {
        if (this.node == null) {
            return null;
        }
        final PeerId leader = this.node.getLeaderId();
        if (leader == null || leader.isEmpty()) {
            return null;
        }
        return leader.getEndpoint();
    }

    /**
     * 【新增】获取此 Raft 节点关联的状态机实例。
     * 主要用于集成测试和诊断，以查询状态机的内部状态（例如，是否正在恢复中）。
     *
     * @return SnackMQStateMachine 实例
     */
    public SnackMQStateMachine getStateMachine() {
        // 【最终修复】: RaftGroupService 和 Node 接口都没有直接的 getFsm() 方法。
        // 正确的做法是通过 Node 实例获取其 NodeOptions，再从 NodeOptions 中获取状态机。
        // Node.getOptions().getFsm() 返回的是一个通用的 StateMachine，
        // 我们需要将其强制转换为我们自己的具体实现 SnackMQStateMachine，
        // 才能调用 isRecovering() 等自定义方法。
        if (this.node == null) {
            log.warn("Attempted to get state machine before Raft node is started.");
            return null;
        }
        return (SnackMQStateMachine) this.node.getOptions().getFsm();
    }

}